package cn.jinbyte.web.wrapper;

import cn.jinbyte.web.utils.XssUtils;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletRequestWrapper;

import java.util.HashMap;
import java.util.Map;

/**
 * XSS请求包装器，对输入进行HTML编码
 */
public class XssRequestWrapper extends HttpServletRequestWrapper {

    public XssRequestWrapper(HttpServletRequest request) {
        super(request);
    }

    @Override
    public String getParameter(String name) {
        String value = super.getParameter(name);
        return stripXSS(value);
    }

    @Override
    public String[] getParameterValues(String name) {
        String[] values = super.getParameterValues(name);
        if (values != null) {
            int length = values.length;
            String[] encodedValues = new String[length];
            for (int i = 0; i < length; i++) {
                encodedValues[i] = stripXSS(values[i]);
            }
            return encodedValues;
        }
        return super.getParameterValues(name);
    }

    @Override
    public Map<String, String[]> getParameterMap() {
        Map<String, String[]> parameterMap = super.getParameterMap();
        Map<String, String[]> encodedMap = new HashMap<>();

        for (Map.Entry<String, String[]> entry : parameterMap.entrySet()) {
            String[] values = entry.getValue();
            String[] encodedValues = new String[values.length];

            for (int i = 0; i < values.length; i++) {
                encodedValues[i] = stripXSS(values[i]);
            }
            encodedMap.put(entry.getKey(), encodedValues);
        }
        return encodedMap;
    }

    @Override
    public String getHeader(String name) {
        String value = super.getHeader(name);
        return stripXSS(value);
    }

    /**
     * 使用OWASP编码器处理XSS
     */
    private String stripXSS(String value) {
        String result = value;
        if (value != null && !value.isEmpty()) {
            result = XssUtils.sanitizeGeneral(value);
        }
        // todo 处理

        return result;
    }
}
